from typing import Optional, List, Dict, Any
import litellm
from litellm import completion
import os
from dotenv import load_dotenv
import json

# Load environment variables
load_dotenv()

class LLMService:
    def __init__(self):
        # Set API keys from environment variables
        huggingface_api_key = os.getenv("HUGGINGFACE_API_KEY")
        openai_api_key = os.getenv("OPENAI_API_KEY")
        anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
        ollama_base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
        
        # Configure
        if huggingface_api_key:
            os.environ["HUGGINGFACE_API_KEY"] = huggingface_api_key
        if openai_api_key:
            litellm.openai_key = openai_api_key
        if anthropic_api_key:
            litellm.anthropic_key = anthropic_api_key
            
        # Set Ollama base URL
        litellm.api_base = ollama_base_url
        
        # Default model
        self.default_model = "huggingface/microsoft/DialoGPT-large"
        
    def generate_text(
        self, 
        prompt: str, 
        model: Optional[str] = None, 
        max_tokens: int = 1000,
        temperature: float = 0.7
    ) -> str:
        """
        Generate text using the specified LLM
        """
        if model is None:
            model = self.default_model
            
        try:
            response = completion(
                model=model,
                messages=[{"content": prompt, "role": "user"}],
                max_tokens=max_tokens,
                temperature=temperature
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"Error generating text: {str(e)}")
            raise e

    def generate_video_script(self, topic: str, style: str = "narrative", duration: int = 60) -> Dict[str, Any]:
        """
        Generate a video script based on topic and style
        """
        prompt = f"""
        Create a video script about "{topic}" in a {style} style.
        The video should be approximately {duration} seconds long.
        Include:
        1. A compelling introduction
        2. Main content with clear points
        3. A conclusion
        4. Suggested visual elements for each section
        5. Estimated timing for each section
        
        Format the response as JSON with the following structure:
        {{
            "title": "...",
            "introduction": {{
                "text": "...",
                "visual_suggestions": "...",
                "duration": ...
            }},
            "sections": [
                {{
                    "title": "...",
                    "text": "...",
                    "visual_suggestions": "...",
                    "duration": ...
                }}
            ],
            "conclusion": {{
                "text": "...",
                "visual_suggestions": "...",
                "duration": ...
            }},
            "total_duration": ...
        }}
        """
        
        try:
            response_text = self.generate_text(prompt, max_tokens=2000)
            # Extract JSON from response if it contains other text
            start_idx = response_text.find('{')
            end_idx = response_text.rfind('}') + 1
            if start_idx != -1 and end_idx != 0:
                json_str = response_text[start_idx:end_idx]
                return json.loads(json_str)
            else:
                # If no JSON found, return the plain text
                return {"script": response_text}
        except Exception as e:
            print(f"Error generating video script: {str(e)}")
            return {"error": str(e)}

    def generate_video_ideas(self, topic: str, count: int = 5) -> List[str]:
        """
        Generate video ideas based on a topic
        """
        prompt = f"""
        Generate {count} creative video ideas about "{topic}".
        Each idea should be unique and engaging.
        Focus on concepts that could be turned into interesting videos.
        """
        
        response = self.generate_text(prompt, max_tokens=1000)
        ideas = [idea.strip() for idea in response.split('\n') if idea.strip() and idea[0].isdigit()]
        return ideas if ideas else [response]

    def improve_prompt(self, prompt: str) -> str:
        """
        Improve a text prompt for better AI video generation
        """
        improvement_prompt = f"""
        Improve this prompt for AI video generation:
        "{prompt}"
        
        Provide a more detailed, specific, and vivid description that would help 
        AI models generate better video content. Include details about:
        - Visual elements
        - Style and mood
        - Color palette
        - Camera movements
        - Any specific elements to include
        
        Return only the improved prompt.
        """
        
        improved = self.generate_text(improvement_prompt, max_tokens=500)
        return improved.strip()
